


import time
import numpy as np
import os
from functools import partial
import jax.numpy as jnp
import jax
import orbax.checkpoint as orbax
import optax
from flax import nnx
from tqdm import trange


os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"


JOINT_MAX = np.array([150, 180, 0, 100, 70, 120])
JOINT_MIN = np.array([-150, 0, -170, -100, -70, -120])

# def denormalize_joint(joints):
#     return (np.array(joints) + 1)/2*(JOINT_MAX - JOINT_MIN) + JOINT_MIN

class EndJointModel(object):
    '''
    Model:
    map arm locations (x, y, z, rx, ry, rz) to (j1, j2, j3, j4, j5, j6)
    '''
    def __init__(self,
                 model: nnx.Module, 
                 max_grad_norm=1,
                 lr=1e-3,
                 model_file: str = None):
        # define the model
        self.model = model
        opt = optax.chain(optax.clip_by_global_norm(max_norm=max_grad_norm), optax.adamw(learning_rate=lr))
        self.opt = nnx.Optimizer(self.model, opt)
        self.ema_params = nnx.Param(model)  #  nnx.update(nnx.Param(model), ema_params)
        if model_file is not None:
            self.load_ckpt(model_file)
    
    def __call__(self, gripper_pos):
        if gripper_pos.ndim == 1:
            # add batch dim
            gripper_pos = gripper_pos[None, :]
        return self._jit_call(self.model, gripper_pos)
        # return denormalize_joint(self.network.apply_fn(self.network.params, gripper_pos))

    @partial(nnx.jit, static_argnums=(0,)) 
    def _jit_call(self, model: nnx.Module, x):
        return model(x)  # jit compiled so that inference can by very fast!
    
    def update(self, batch):
        info = self.jit_train_step(self.model, self.opt, batch)
        return info
    
    @staticmethod
    def l2_loss_fn(model: nnx.module, batch):
        pred_joints = model(batch.x) 
        loss = jnp.mean((pred_joints - batch.y)**2)
        return loss, {"loss": loss}
        
    @partial(nnx.jit, static_argnums=(0,))
    def jit_train_step(self, model: nnx.Module, optimizer: nnx.Optimizer, batch):
        '''
        batch.arm: (batch_size, 6d)
        batch.joint: (batch_size, 6d)
        '''
        (_, info), grads = nnx.value_and_grad(self.l2_loss_fn, has_aux=True)(model, batch)
        optimizer.update(grads)  # inplace update
        return info
    
    def train(self, dataloader, num_epochs: int, save_as: str):
        from tensorboardX import SummaryWriter
        writer = SummaryWriter(log_dir=f"log/arm_joint/{save_as}_{round(time.time())}")
        
        avg_loss=0
        
        for epoch in trange(num_epochs, desc="Train"):
            
            cnt_batch = 0
            loss_rec = []
            
            for batch in dataloader:
                info = self.update(batch)
                loss_rec.append(info["loss"].item())
                cnt_batch += 1
                    
            avg_loss = sum(loss_rec)/len(loss_rec)
            writer.add_scalar("train/loss", avg_loss, epoch)
            
            if (1 + epoch) % 10 == 0:
                print(f"Epoch{epoch+1}: loss = {avg_loss}")
                
            if (1 + epoch) % 100 == 0:
                self.save_ckpt(file_name=save_as)
            
        self.save_ckpt(file_name=save_as)

    def save_ckpt(self, file_name: str, save_dir: str="data/weights"):
        
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        
        # Checkpoint path should be absolute
        filename = os.path.join(os.getcwd(), save_dir, file_name)
        # Save the parameters
        checkpointer = orbax.PyTreeCheckpointer()
        checkpointer.save(filename, nnx.state(self.model), force=True)
        print(f"Successfully save model to {filename}")

    def load_ckpt(self, file_name: str, save_dir: str="data/weights"):
        
        filename = os.path.join(os.getcwd(), save_dir, file_name)
        state = nnx.state(self.model)
        # Load the parameters
        checkpointer = orbax.PyTreeCheckpointer()
        state = checkpointer.restore(filename, item=state)
        # update the model with the loaded state
        nnx.update(self.model, state)
        print(f"Successfully load model from {filename}")
        return self.model




# import time
# import logging
# import numpy as np
# import json
# import os
# from functools import partial
# import jax.numpy as jnp
# import jax
# import orbax.checkpoint
# import optax
# from flax import nnx
# from flax.training.train_state import TrainState
# from collections import namedtuple
# from model.utils.jax_types import Params
# from model.networks.mlp import MLP
# from tqdm import trange
# import torch
# from torch.utils.data import Dataset
# from data.jax_dataloader import NumpyLoader


# os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
# Batch = namedtuple("Batch", "gripper joint")


# class EndJointModel(object):
#     '''
#     Model:
#     map gripper locations (x, y, z) to (j1, j2, j3, j4, j5, j6)
#     '''
#     def __init__(self, 
#                  seed=0,
#                  max_grad_norm=1,
#                  lr=1e-3,
#                  model_file: str = None):
#         self.rng = jax.random.PRNGKey(seed)
#         self.rng, key = jax.random.split(self.rng)
        
#         # define the model
#         network = MLP(output_dim=6, 
#                       activation_type="Mish",
#                       hidden_dims=(512, 512, 512),  # (512, 512, 512),
#                       use_layernorm = True)
#         params = network.init( key,  jnp.empty((1, 3)))

#         self.network = TrainState.create(
#             apply_fn=network.apply,
#             params=params, 
#             tx=optax.chain(optax.clip_by_global_norm(max_norm=max_grad_norm), optax.adamw(learning_rate=lr)),
#         )
        
#         if model_file is not None:
#             self.load_ckpt(model_file)
            
#     def __call__(self, gripper_pos):
#         if gripper_pos.ndim == 1:
#             # add batch dim
#             gripper_pos = gripper_pos[None, :]
#         return self.network.apply_fn(self.network.params, gripper_pos)
#         # return denormalize_joint(self.network.apply_fn(self.network.params, gripper_pos))
    
#     def update(self, batch: Batch):
#         self.network, info = self._jit_update(self.network, batch)
#         return info
            
#     @partial(jax.jit, static_argnums=(0,))
#     def _jit_update(self, model: TrainState, batch: Batch):
#         '''
#         batch.gripper_pos: (batch_size, 6d)
#         batch.joint_pos: (batch_size, 6d)
#         '''
#         def loss_fn(params: Params):
#             pred_joints = model.apply_fn(params, batch.gripper)  # xyz
#             loss = jnp.mean((pred_joints - batch.joint)**2)
            
#             return loss, {"loss": loss}
        
#         (_, info), grads = jax.value_and_grad(loss_fn, has_aux=True)(model.params)
#         new_model = model.apply_gradients(grads=grads)
#         return new_model, info
    
#     @partial(jax.jit, static_argnums=(0,))
#     def _jit_pred_joints(self,params: Params, gripper_pos: jnp.ndarray):
#         return self.network.apply_fn(params, gripper_pos)
    
#     def train(self, dataloader, num_epochs: int, save_as: str):
#         from tensorboardX import SummaryWriter
#         writer = SummaryWriter(log_dir=f"log/end_joint/{save_as}_{round(time.time())}")
        
#         avg_loss=0
        
#         for epoch in trange(num_epochs, desc="Train"):
            
#             cnt_batch = 0
#             loss_rec = []
            
#             for batch in dataloader:
#                 info = self.update(batch)
#                 loss_rec.append(info["loss"].item())
#                 cnt_batch += 1
                    
#             avg_loss = sum(loss_rec)/len(loss_rec)
#             writer.add_scalar("train/loss", avg_loss, epoch)
            
#             if (1 + epoch) % 100 == 0:
#                 print(f"Epoch{epoch+1}: loss = {avg_loss}")
#                 self.save_ckpt(file_name=save_as)
            
#         self.save_ckpt(file_name=save_as)

#     def save_ckpt(self, file_name: str, save_dir: str="data/weights"):
        
#         if not os.path.exists(save_dir):
#             os.mkdir(save_dir)
            
#         filename = os.path.join(os.getcwd(), save_dir, file_name)

#         orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
#         orbax_checkpointer.save(filename, {"params": self.network.params}, force=True)
#         print(f"Successfully save model to {filename}")

#     def load_ckpt(self, file_path: str="data/weights/gp_mapping.ckp"):
#         orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
#         filename = os.path.join(os.getcwd(), file_path)
#         raw_restored = orbax_checkpointer.restore(filename)  
#         self.network = self.network.replace(params=raw_restored["params"])
    
#         print(f"Successfully load model from {filename}")

